{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Iterable 与 Iterator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在我们开始学习数据容器之前，先来学习 *iterable* 和 *iterator* 的概念。\n",
    "\n",
    "理解这两个概念就能理解 `for...in` 循环的本质，同时 *iterable* 也是所有数据容器的共性特征。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `for...in` 循环的本质"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Iterable* 和 *iterator* 这两个词，来源于动词 *iterate*，*iterate* 与其名词形式 *iteration* 的含义就是“重复”，重复做一件事或者重复一句话；在计算机领域通常翻译为“迭代”，对一组数据逐个执行特定操作都可以称为 *iteration*，比如你已经熟悉的 `for...in` 循环就是例子："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "2\n",
      "3\n"
     ]
    }
   ],
   "source": [
    "for i in [1, 2, 3]:\n",
    "    print(i)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "上面的代码本质是这样：\n",
    "1. 取列表 `[1, 2, 3]` 中的第一个元素（`1`），令 `i = 1`；\n",
    "2. 执行 `print(i)`，打印出 `i` 的值；\n",
    "3. 取列表 `[1, 2, 3]` 中的“下一个”元素，令 `i` 的值等于该元素；\n",
    "4. 重复步骤 2 至 3，直到列表中没有元素（即“下一个”元素为空值）。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这个概念抽象一下，迭代的本质是对一组数据执行特定操作，确保所有数据都被操作正好一次，不重不漏。这是非常常用的一种场景，这里面有两个要素：\n",
    "* 被迭代处理的数据集 `X`，其中包含若干元素，可以通过“给我下一个元素”这样的指令逐个取出这些元素，确保不重不漏；\n",
    "* 需要对 X 的每个元素执行的操作 `f()`。\n",
    "\n",
    "只要具备这两个要素，我们就可以通过下面这样的代码来遍历操作 X 中的所有元素：\n",
    "```python\n",
    "for x in X:\n",
    "    f(x)\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Python 中的 *iterable* 和 *iterator* 就是从上面描述的抽象模型里发展而来的设计，给出了在 Python 中做迭代的标准模式。\n",
    "\n",
    "* 一个对象 X 是“**可迭代的**（*iterable*）”，就是说这个对象支持一个名为 `__iter__()` 的方法，这个方法返回一种叫“**迭代器**（*iterator*）”的对象；\n",
    "* 一个**迭代器**（*ietrator*）必须支持一个名为 `__next__()` 的方法，这个方法每次返回 X 里的下一个元素，直到所有元素被遍历正好一次。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "而我们用了好多次的 `for...in` 循环实际上是这么实现的：\n",
    "1. 获得 X 的迭代器 `iter = X.__iter__()`；\n",
    "2. 获得 X 中下一个元素 `x = iter.__next__()`；\n",
    "3. 如果 x 不为空，则执行 `f(x)` 然后重复步骤 2~3；否则结束。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这下我们终于可以揭晓了：**凡是 *iterable* 的东西，都可以放在 `for...in` 循环中进行迭代操作**。\n",
    "\n",
    "我们后面要介绍的数据容器（*list*、*tuple*、*dictionary*、*set*）都是 *iterable*，所以都可以用 `for-in` 来遍历。\n",
    "\n",
    "另外，X 可以同时支持 `__iter__()` 和 `__next__()`，这样 X 既是 *iterable* 也是自己的 *iterator*，可以省略取迭代器的步骤，直接用 `__next__()` 开始，上面列出的几种内置数据容器都是这么做的。\n",
    "\n",
    "下面我们创建一些自己的 *iterable* 试试。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 迭代器例一：平方数序列"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "第一个例子我们返回自然数的平方数序列，也就是 `1, 4, 9, 16, 25, ...`\n",
    "\n",
    "*Iterator* 是一种对象，所以首先我们要使用 `class` 关键字来定义一个类："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class squares:\n",
    "    # 实例变量 self.base 用来记录当前平方数的基数，这个基数应该从 1 开始逐渐递增。\n",
    "    # 然后定义了三个标准方法：\n",
    "    #   __init__() 把 self.base 设置为初始值 0；\n",
    "    #   __iter__() 返回迭代器，这里就是自己，我们直接返回 self；\n",
    "    #   __next__() 是迭代器主方法，这里我们先把 self.base 加一，然后返回它的平方，由于初始为 0，第一次调用返回的是 1 的平方。\n",
    "    # 最后定义 next 作为 __next__ 的别名，方便调用（也与 Python 2.x 兼容）。\n",
    "    def __init__(self):\n",
    "        self.base = 0\n",
    "        \n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        self.base += 1\n",
    "        return self.base * self.base\n",
    "\n",
    "    next = __next__\n",
    "\n",
    "# 创建 squares 类的实例对象 iter，然后就可以调用其 next() 方法来迭代\n",
    "iter = squares()\n",
    "iter.next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iter.next()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "iter.next()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "注意到我们上面没有任何约束条件，这个迭代可以无限进行下去，如果我们把这个迭代器用在 `for...in` 循环中的话就会出现无限循环（直到计算机耗尽内存资源），所以通常我们会希望有个约束条件让迭代到一定程度就停下来，我们可以改写一下上面的代码："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class squares:\n",
    "    # 这次我们的初始化方法增加了一个参数 high\n",
    "    # 初始化时将其保存在一个新增实例变量 self.high 中，限制迭代时 self.base 不能超过这个数\n",
    "    def __init__(self, high):\n",
    "        self.base = 0\n",
    "        self.high = high\n",
    "        \n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    # 在迭代器主方法里加入判断，如果 self.base 超过了 self.high 就扔出一个 StopIteration 异常\n",
    "    def __next__(self):\n",
    "        self.base += 1\n",
    "        if self.base > self.high:\n",
    "            raise StopIteration\n",
    "        else:\n",
    "            return self.base * self.base\n",
    "\n",
    "    next = __next__\n",
    "\n",
    "# 这次我们实例化时要传入一个参数来指定 self.high 的值\n",
    "# 这将使得我们调用三次 next()，输出 1，4，9 之后再调用就不会有返回值\n",
    "iter = squares(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "这次我们的迭代器有“边界”了，就可以将其直接用在 `for...in` 循环中，`for...in` 循环会在碰到 `StopIteration` 时自动终止："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "4\n",
      "9\n"
     ]
    }
   ],
   "source": [
    "for n in squares(3):\n",
    "    print(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 迭代器例二：斐波那契数列"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "第二个例子我们来实现一个**斐波那契**（*Fibonacci*）数列，这个数列之前已经见过，其特点是每一个元素它前面两个元素之和，写下来大致是这个样子的：`1, 1, 2, 3, 5, 8, 13, 21, ...`\n",
    "\n",
    "基本的迭代器实现和例一没啥区别，如果前面的内容没有什么问题，这个也就很好理解："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class fib:\n",
    "    def __init__(self):\n",
    "        self.prev = 0\n",
    "        self.curr = 1\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        self.prev, self.curr = self.curr, self.prev + self.curr\n",
    "        return self.prev"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "可以看到我们这个实现也是可以无限迭代的，如果我们需要限制迭代次数，可以像例一那样改写，但是大多数情况下不需要，因为这种操作太常见了，早就有了标准化的通用解决方案——这也是学编程的一个窍门：凡是看上去很“常规”和“常见”的需求，一般都会有优秀的通用解决方案，我们找出来用比自己写好，这叫“**不重复发明轮子**（*Don't Reinvent The Wheel*）”原理。这些标准解决方案在 `itertools` 模块中，用于对无限迭代器进行切片的方法叫做 `islice()`。\n",
    "\n",
    "`islice()` 做的事情相当于从迭代器产生的无限序列中切出一段来，它返回的也是一个迭代器，返回的迭代器输出的就是有头有尾、有限的序列了。\n",
    "\n",
    "`islice()` 接受三个参数：一个迭代器实例，切片的起始和结束序号（序号从 0 开始算）；切的时候和 `range()` 函数的算法一样，包含起始序号那一项，但不包括结束序号的一项。比如 `islice(f, 0, 10)` 就是从迭代器 `f` 生成的序列中切出第 1 到第 10 项，返回这 10 项对应的有限迭代器。\n",
    "\n",
    "其实 `islice()` 还可以有第四个参数：步长 `step`，缺省值为 1，就是起始到结束范围内每一项都保留；如果指定了别的 `step` 值，那么就是每隔 `step-1` 项输出一项。基本逻辑和 `range()` 函数一样。\n",
    "\n",
    "因为 `islice()` 返回的还是迭代器，所以可以直接用于 `for...in` 循环中。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1\n",
      "1\n",
      "2\n",
      "3\n",
      "5\n",
      "8\n",
      "13\n",
      "21\n",
      "34\n",
      "55\n"
     ]
    }
   ],
   "source": [
    "from itertools import islice\n",
    "\n",
    "f = fib()\n",
    "for n in islice(f, 0, 10):\n",
    "    print(n)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`islice()` 帮我们做了限界这样的事情，我们在写迭代器的时候就不用考虑了，只要把迭代的算法写清楚，要截取一段来用的话用 `islice()` 就好了。`itertools` 里还有不少很有用的东西，比如 `cycle()` 可以循环一个有限的迭代器得到一个无限版本："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['red', 'yellow', 'green', 'red', 'yellow', 'green', 'red']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from itertools import cycle\n",
    "# 数组既是 iterable 也是 iterator，cycle() 从输入的有限数组循环创建一个无限的 iterator\n",
    "colors = cycle(['red', 'yellow', 'green'])\n",
    "# 用 islice() 截取这个无限迭代器的前 10 个，然后用 list() 把输出的迭代器转换成列表\n",
    "list(islice(colors, 0, 7))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "更多关于 `itertools` 库的说明可以参考 [官方文档](https://docs.python.org/3.7/library/itertools.html)。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## One More Thing..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "不知道你有没有在书写上面的代码时觉得每次要定义个 `class` 还有 `__iter__()` `__next__()` 这样的“八股”很啰嗦？恭喜你，很有优秀程序员的潜质，事实上上面代码里真正有价值的就是 `__next__()` 的实现部分，所以 Python 提供了一种特殊的迭代器（*iterator* ）写法，叫做“**生成器**（*generator*）”，更紧凑更精简。比如上面的 Fibonacci 数列，用生成器来写就是这样的："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[1, 1, 2, 3, 5, 8, 13, 21, 34, 55]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def fib():\n",
    "    prev, curr = 0, 1\n",
    "    while True:\n",
    "        # 这个 yield 是一切魔法的核心，下面解释\n",
    "        yield curr\n",
    "        prev, curr = curr, prev + curr\n",
    "\n",
    "f = fib()\n",
    "list(islice(f, 0, 10))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "generator"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(fib())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "是不是清爽漂亮了很多？但是这个魔法的原理是啥呢？首先我们来看看 `fib()` 这个定义：\n",
    "* 一个 *generator* 定义就是一个函数的定义，但这是个特殊的函数，函数体里没有 `return`，取而代之的是特殊的关键字 `yield`；\n",
    "* 解释器会自动把任何包含 `yield` 的函数包装成 `generator`（一种特殊的 *iterator*）类型的对象，并将此对象作为函数的返回值，所以上面的 `f` 变量其实还是个 *iterator*。\n",
    "\n",
    "然后就是 `yield` 这个魔法的关键了。`yield` 和 `return` 相当，也会从函数中返回一些值，但不一样的是：`return` 返回值的同时彻底终止了函数，而 `yield` 则是“暂停”了函数执行，保存了函数当时的所有状态（比如所有局部变量的值），以后再返回到 `yield` 的下一行继续执行——这个“以后”到底是什么时候呢？聪明的你一定可以猜到，就是我们下一次调用 `__next__()` 的时候！\n",
    "\n",
    "小结一下，上面的代码定义了一个函数 `fib()`，调用这个函数会返回一个生成器（一种特殊的迭代器）：\n",
    "* 当我们第一次调用这个生成器的 `__next__()` 的时候，会执行 `fib()` 函数的函数体直到碰到 `yield` 语句，`yield` 的返回值就是这次调用 `__next__()` 的结果；同时解释器暂停函数 `fib()` 的运行并把状态保存下来；\n",
    "* 当我们再次调用生成器的 `__next__()` 的时候，解释器重新唤醒函数 `fib()`，恢复保存的状态，并从上次 `yield` 的下一行继续执行，直到再碰到 `yield`，`yield` 的返回值就是这次调用 `__next__()` 的结果；\n",
    "* 每次调用生成器的 `__next__()` 时重复上述操作。\n",
    "\n",
    "是不是略有点烧脑？在这里你一定要烧清楚。这次清楚了后面就简单，否则这一类的逻辑还会有很多，始终是要迈过的坎，而且生成器很有用，很多地方都会碰到。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 练习"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "尝试定义可以生成素数的迭代器和生成器，并利用之输出前 10 个素数。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 参考答案"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "我们先来看迭代器的写法，用到了我们之前写过的 `is_prime()` 算法："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 3, 5, 7, 11, 13, 17, 19, 23, 29]"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from math import sqrt\n",
    "from itertools import islice\n",
    "\n",
    "class primes:\n",
    "    def __init__(self):\n",
    "        self.next = 2\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def _is_prime(self, n):\n",
    "        if n < 2:\n",
    "            return False\n",
    "\n",
    "        if n in (2, 3):\n",
    "            return True\n",
    "\n",
    "        for i in range(2, int(sqrt(n)) + 1):\n",
    "            if n % i == 0:\n",
    "                return False\n",
    "\n",
    "        return True\n",
    "\n",
    "    def __next__(self):\n",
    "        n = self.next\n",
    "        while not self._is_prime(n):\n",
    "            n += 1\n",
    "\n",
    "        self.next = n + 1\n",
    "        return n\n",
    "    \n",
    "list(islice(primes(), 0, 10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "参照前面的例子，将上面这个 `iterator` 类改写为 `generator` 应该也不难，我们把这个作业留给你。\n",
    "\n",
    "下面额外给出一种 *generator* 的写法，这个写法用到了著名的**埃拉托斯特尼筛法**（*Sieve Of Eratosthenes*），是快速计算素数的经典算法，关于筛法的基本介绍可以参考[维基百科](https://zh.wikipedia.org/zh-cn/%E5%9F%83%E6%8B%89%E6%89%98%E6%96%AF%E7%89%B9%E5%B0%BC%E7%AD%9B%E6%B3%95)。\n",
    "\n",
    "我们在代码的每一步都写了简要注释，主要戏法在函数体里定义的那个 `D`，这是一个 `dictionary` 类型的数据容器，如果看不明白也没关系，在学完后面几种数据容器之后再来看，就应该容易一些了。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def primes():\n",
    "    \"\"\"Generator which yields the sequence of prime numbers via the Sieve of Eratosthenes.\"\"\"\n",
    "    D = {}  # map composite integers to primes witnessing their compositeness\n",
    "    q = 2   # first integer to test for primality\n",
    "    while 1:\n",
    "        if q not in D:\n",
    "            yield q        # not marked composite, must be prime\n",
    "            D[q*q] = [q]   # first multiple of q not already marked\n",
    "        else:\n",
    "            for p in D[q]: # move each witness to its next multiple\n",
    "                D.setdefault(p+q,[]).append(p)\n",
    "            del D[q]       # no longer need D[q], free memory\n",
    "        q += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from itertools import islice\n",
    "list(islice(primes(), 0, 20))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
